import numpy as np

# 定义计算两点间距离的函数
def distance(p1, p2):
    """计算两点之间的欧几里得距离"""
    return np.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)


def compute_single_code(preds, gts, threshold):
    # 对真实关键点（gts）进行重排，以便后续处理
    rearranged_gts = np.zeros_like(gts)
    k = 0
    for i in range(4):
        for j in range(i, len(gts), 4):
            rearranged_gts[k] = gts[j]
            k += 1

    # 用于存储匹配成功的结果
    matched_results = []

    # 遍历每组预测点
    for i, group_predicted in enumerate(preds):
        # 根据索引选择对应的真实关键点组
        group_true = rearranged_gts[i * len(gts) // 4:(i + 1) * len(gts) // 4]

        new_group_true = np.array([item[:2] for item in group_true if item[2] == 1])

        # print(group_true)
        # print(group_predicted)

        # 对每个预测点，找出所有距离小于阈值的真实点作为候选匹配点
        candidates = {}
        for p in group_predicted:
            p_tuple = tuple(p[:2])  # 只取前两个元素并转换为元组
            candidates[p_tuple] = [(t, distance(p[:2], t)) for t in new_group_true if distance(p[:2], t) <= threshold]

        # 存储最终的匹配点
        final_matches = {}
        for point, matches in candidates.items():
            for match, dist in matches:
                match_key = match.tobytes()
                if match_key not in final_matches or final_matches[match_key][1] > dist:
                    final_matches[match_key] = (point, dist)

        # 将匹配成功的点添加到结果列表中
        for match, (point, dist) in final_matches.items():
            matched_results.append(point)

    # 返回匹配成功的点的数量
    return len(matched_results)
